
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include "../iface/debug.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "RSAPublicKey.h"
#include "RSAPrivateKey.h"

RSAPrivateKey::RSAPrivateKey()
{
   key = NULL;
   keyname = NULL;
}

RSAPrivateKey::RSAPrivateKey(const char *name)
{
   key = NULL;
   keyname = NULL;
   SetKeyName(name);
}

// Construct a duplicate of the private key
RSAPrivateKey::RSAPrivateKey(const RSAPrivateKey &other)
{
   key = RSA_new();
   if (key == NULL)
   {
      errprint("RSA_new failed");
      return;
   }
   CopyRSASSLPrivateKey(key, other.key);
   keyname = NULL;
   SetKeyName(other.keyname);
}

// Construct from PEM format
RSAPrivateKey::RSAPrivateKey(const char *pemtext, const char *name)
{
   ImportKeyFromASCII(pemtext, name);
}

RSAPrivateKey::RSAPrivateKey(const char *pemtext, const char *password, const char *name)
{
   ImportKeyFromASCIIWithPassword(pemtext, password, name);
}

// Construct from OpenSSL format
RSAPrivateKey::RSAPrivateKey(RSA *prvkey, const char *name)
{
   key = RSA_new();
   if (key == NULL)
   {
      errprint("RSA_new failed");
      return;
   }
   CopyRSASSLPrivateKey(key, prvkey);
   keyname = NULL;
   SetKeyName(name);
}

RSAPrivateKey::~RSAPrivateKey()
{
   if (key) RSA_free(key);
   key = NULL;
   if (keyname) free((void *)keyname);
   keyname = NULL;
}

// Get context
int RSAPrivateKey::GetCryptlibContext(CRYPT_CONTEXT *keyContext)
{
   int status;
   CRYPT_PKCINFO_RSA *rsaKey = (CRYPT_PKCINFO_RSA *)malloc(sizeof(CRYPT_PKCINFO_RSA));
   if (GetCryptlibKeyComponents(rsaKey) != 0)
   {
      free(rsaKey);
      return -1;
   }
   int magictriescntxt = 0;
   do
   {
      status = cryptCreateContext(keyContext, CRYPT_UNUSED, CRYPT_ALGO_RSA);
      if (status == CRYPT_ERROR_TIMEOUT) apr_sleep(rand() % CRYPT_MAX_TIMESLEEP);
      magictriescntxt++;
   } while ((status == CRYPT_ERROR_TIMEOUT) && (magictriescntxt < CRYPT_MAX_RETRIES));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keyContext, CRYPT_CTXINFO_LABEL, keyname, strlen(keyname));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keyContext, CRYPT_CTXINFO_KEY_COMPONENTS, rsaKey, sizeof(CRYPT_PKCINFO_RSA));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   cryptDestroyComponents(rsaKey);
   free(rsaKey);
   return 0;
failover:
   cryptDestroyComponents(rsaKey);
   free(rsaKey);
   return -1;
}

// Get context
int RSAPrivateKey::GetCryptlibDecryptionContext(CRYPT_CONTEXT *keyContext)
{
   return GetCryptlibContext(keyContext);
}

// Get context
int RSAPrivateKey::GetCryptlibSigningContext(CRYPT_CONTEXT *keyContext)
{
   return GetCryptlibContext(keyContext);
}

// Get key components
int RSAPrivateKey::GetCryptlibKeyComponents(CRYPT_PKCINFO_RSA *rsaKey)
{
   unsigned char *buffera[8];
   cryptInitComponents(rsaKey, CRYPT_KEYTYPE_PRIVATE);
   buffera[0] = (unsigned char *)malloc(BN_num_bytes(key->n));
   if (!(buffera[0])) { errprint("Could not allocate buffer for key->n."); return -1; }
   buffera[1] = (unsigned char *)malloc(BN_num_bytes(key->e));
   if (!(buffera[1])) { errprint("Could not allocate buffer for key->e."); return -1; }
   buffera[2] = (unsigned char *)malloc(BN_num_bytes(key->d));
   if (!(buffera[2])) { errprint("Could not allocate buffer for key->d."); return -1; }
   buffera[3] = (unsigned char *)malloc(BN_num_bytes(key->p));
   if (!(buffera[3])) { errprint("Could not allocate buffer for key->p."); return -1; }
   buffera[4] = (unsigned char *)malloc(BN_num_bytes(key->q));
   if (!(buffera[4])) { errprint("Could not allocate buffer for key->q."); return -1; }
   buffera[5] = (unsigned char *)malloc(BN_num_bytes(key->iqmp));
   if (!(buffera[5])) { errprint("Could not allocate buffer for key->iqmp."); return -1; }
   buffera[6] = (unsigned char *)malloc(BN_num_bytes(key->dmp1));
   if (!(buffera[6])) { errprint("Could not allocate buffer for key->dmp1."); return -1; }
   buffera[7] = (unsigned char *)malloc(BN_num_bytes(key->dmq1));
   if (!(buffera[7])) { errprint("Could not allocate buffer for key->dmq1."); return -1; }
   BN_bn2bin(key->n, buffera[0]);
   BN_bn2bin(key->e, buffera[1]);
   BN_bn2bin(key->d, buffera[2]);
   BN_bn2bin(key->p, buffera[3]);
   BN_bn2bin(key->q, buffera[4]);
   BN_bn2bin(key->iqmp, buffera[5]);
   BN_bn2bin(key->dmp1, buffera[6]);
   BN_bn2bin(key->dmq1, buffera[7]);
   cryptSetComponent(rsaKey->n, buffera[0], BN_num_bits(key->n));
   cryptSetComponent(rsaKey->e, buffera[1], BN_num_bits(key->e));
   cryptSetComponent(rsaKey->d, buffera[2], BN_num_bits(key->d));
   cryptSetComponent(rsaKey->p, buffera[3], BN_num_bits(key->p));
   cryptSetComponent(rsaKey->q, buffera[4], BN_num_bits(key->q));
   cryptSetComponent(rsaKey->u, buffera[5], BN_num_bits(key->iqmp));
   cryptSetComponent(rsaKey->e1, buffera[6], BN_num_bits(key->dmp1));
   cryptSetComponent(rsaKey->e2, buffera[7], BN_num_bits(key->dmq1));
   for (int i = 0; i < 8; i++) free(buffera[i]);
   return 0;
}

// Get new OpenSSL RSA object
int RSAPrivateKey::GetSSLKeyComponents(RSA **prvkey)
{
   *prvkey = RSA_new();
   if ((*prvkey) == NULL)
   {
      errprint("RSA_new failed");
      return -1;
   }
   return CopyRSASSLPrivateKey(*prvkey, key);
}

// Get PEM data
int RSAPrivateKey::GetPEMKeyData(char **pemtext)
{
   char *data;
   int damt;
   int retval = 0;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(epkey, key);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PrivateKey(bo, epkey, 0, 0, 0, 0, 0))
   {
      retval = 1;
      errprint("could not export private key");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   *pemtext = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(*pemtext, data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
//   printf("%s\n", dbuf);
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

// Get PEM key data
int RSAPrivateKey::GetPEMKeyData(char **pemtext, const char *password)
{
   char *data;
   int damt;
   int retval = 0;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_RSA(epkey, key);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PKCS8PrivateKey(bo, epkey, EVP_des_ede3_cbc(), NULL, 0, 0, (void *)password))
   {
      retval = 1;
      errprint("could not export private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   (*pemtext) = (char *)malloc(sizeof(char)*(damt+1));
   memcpy((*pemtext), data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

// Compare key to key
int RSAPrivateKey::CompareKeyBelongs(PublicKey *pubkey)
{
   if (pubkey == NULL)
   {
      errprint("comparing against NULL key");
      return 0;
   }
   switch (pubkey->GetKeyType())
   {
      case KEY_TYPE_PUBLIC_RSA:
      {
         RSA *rsaKeyOther = NULL;
         ((RSAPublicKey *)pubkey)->GetSSLKeyComponents(&rsaKeyOther);
         EVP_PKEY *evp1 = EVP_PKEY_new();
         EVP_PKEY *evp2 = EVP_PKEY_new();
         EVP_PKEY_set1_RSA(evp1, key);
         EVP_PKEY_set1_RSA(evp2, rsaKeyOther);
         int retval = EVP_PKEY_cmp(evp1, evp2);
         EVP_PKEY_free(evp1);
         EVP_PKEY_free(evp2);
         RSA_free(rsaKeyOther);
         fSSLCheckError();
         switch (retval)
         {
            case -2:
               errprint("key compare operation not supported (EVP_PKEY_cmp())");
               break;
            case -1:
               errprint("tried to compare different key types (EVP_PKEY_cmp())");
               return 0;
            case 0:
               return 0;
            case 1:
               return 1;
         }
         return 0;
         break;
      }
      default:
      {
         return 0;
         break;
      }
   }
}

int RSAPrivateKey::CopyRSASSLPrivateKey(RSA *dst, RSA *src)
{
   if (dst == NULL) { errprint("NULL private key destination"); return -1; }
   if (src == NULL) { errprint("NULL private key source"); return 1; }
   // Copy public key
   dst->n = BN_new();
   if (BN_copy(dst->n, src->n) != (dst->n)) { errprint("Could not copy prv->n"); return -1; }
   dst->e = BN_new();
   if (BN_copy(dst->e, src->e) != (dst->e)) { errprint("Could not copy prv->e"); return -1; }

   // Copy private key
   dst->d = BN_new();
   if (BN_copy(dst->d, src->d) != (dst->d)) { errprint("Could not copy prv->d"); return -1; }
   dst->p = BN_new();
   if (BN_copy(dst->p, src->p) != (dst->p)) { errprint("Could not copy prv->p"); return -1; }
   dst->q = BN_new();
   if (BN_copy(dst->q, src->q) != (dst->q)) { errprint("Could not copy prv->q"); return -1; }

   dst->dmp1 = BN_new();
   if (BN_copy(dst->dmp1, src->dmp1) != (dst->dmp1)) { errprint("Could not copy prv->dmp1"); return -1; }
   dst->dmq1 = BN_new();
   if (BN_copy(dst->dmq1, src->dmq1) != (dst->dmq1)) { errprint("Could not copy prv->dmq1"); return -1; }
   dst->iqmp = BN_new();
   if (BN_copy(dst->iqmp, src->iqmp) != (dst->iqmp)) { errprint("Could not copy prv->iqmp"); return -1; }
   return 0;   
}

int RSAPrivateKey::ImportPrivateKeyWithOptionalPassword(const char *pemkeydata, const char *password)
{
   int retval = 0;
   BIO *bo = BIO_new_mem_buf((void *)pemkeydata, -1);
   EVP_PKEY *epkey = PEM_read_bio_PrivateKey(bo, NULL, 0, (void *)password);
   if (!epkey)
   {
      retval = 1;
      errprint("could not import private key (password optional)");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   if (key) RSA_free(key);
   key = NULL;
   key = EVP_PKEY_get1_RSA(epkey);
   if (key == NULL)
   {
      retval = -1;
      errprint("failed private key import");
   }
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

int RSAPrivateKey::ImportKeyFromASCII(const char *text, const char *name)
{
   return ImportKeyFromASCIIWithPassword(text, NULL, name);
}

int RSAPrivateKey::ImportKeyFromASCIIWithPassword(const char *text, const char *password, const char *name)
{
   SetKeyName(name);
   int retval = ImportPrivateKeyWithOptionalPassword(text, NULL);
   return retval;
}

int RSAPrivateKey::ExportKeyToASCII(char **text)
{
   return GetPEMKeyData(text);
}

int RSAPrivateKey::ExportKeyToASCIIWithPassword(char **text, const char *password)
{
   return GetPEMKeyData(text, password);
}
